#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import numpy as np
import cv2
from collections import OrderedDict
import glob
import time
import multiprocessing
import torch
import random

from distance import euclidean_squared_distance, euclidean_squared_distance2, cosine_distance, cosine_distance2, cosine_distance_ex

sys.path.append('.')

"""
    过滤掉太相似的人脸
"""


def read_image(image_dir, file_name, img_width, img_height, mode=1):
    file_path = os.path.join(image_dir, file_name.strip())
    img = cv2.imread(file_path)
    if mode == 0:
        img = cv2.resize(img, (img_width, img_height))
    else:
        w, h = img.shape[1], img.shape[0]
        wscale = img_width / w
        hscale = img_height / h
        scale = min(wscale, hscale)
        w = int(round(w * scale))
        h = int(round(h * scale))
        image = cv2.resize(img, (w, h))
        img = np.full((img_height, img_width, image.shape[2]), fill_value=196, dtype=np.uint8)
        x = (img_width-w)//2
        y = (img_height-h)//2
        img[y:y+h, x:x+w, :] = image
    return img


def show_images(image_dir, file_names, dis_text='', win_name='image', wait=False,
                h_image_count=10, v_image_count=10, img_width=96, img_height=96):
    display_image_count = h_image_count * v_image_count
    files = []
    total_image_count = len(file_names)
    last_idx = total_image_count - 1
    # if total_image_count > 0:
    #     print('{}: {}'.format(win_name, file_names[0]))
    for idx, file_name in enumerate(file_names):
        files.append(file_name)
        file_count = len(files)
        if file_count < display_image_count and (idx < last_idx):
            continue
        imgs = []
        for file_name in files:
            if file_name is None:
                img = np.zeros((img_height, img_width, 3), dtype=np.uint8)
            else:
                img = read_image(image_dir, file_name.strip(), img_width, img_height)
            imgs.append(img)
        files.clear()
        image_count = len(imgs)
        images = []
        x_off = 0
        y_off = 0
        for j, img in enumerate(imgs):
            if j % h_image_count == 0:
                x_off = 0
                if j != 0:
                    y_off += img_height
                h_images = [img]
            else:
                x_off += img_width
                h_images.append(img)
            if ((j + 1) % h_image_count == 0) or ((j + 1) == image_count):
                images.append(np.concatenate(h_images, axis=1))

        if len(images) > 1:
            w0 = images[0].shape[1]
            h0 = images[0].shape[0]
            w1 = images[-1].shape[1]
            h1 = images[-1].shape[0]
            if (w0 != w1) or (h0 != h1):
                # image = np.empty_like(images[0])
                image = np.zeros_like(images[0])
                image[0:h1, 0:w1, :] = images[-1]
                images[-1] = image
        draw_image = np.vstack(images)
        cv2.putText(draw_image, '%s %d/%d/%d' % (dis_text, idx + 1, last_idx + 1, total_image_count), (2, 12),
                    cv2.FONT_HERSHEY_PLAIN, 1.0, (0, 255, 0), 1)
        cv2.imshow(win_name, draw_image)
        if wait:
            key = cv2.waitKey()
            if key == 27 or key == ord('q') or key == ord('Q'):
                return True
            elif key == ord('n') or key == ord('N'):
                break
        else:
            break
    return False


def face_select(idxes, distance_thr, distmat):
    same_idx_set = set()
    keep_idxes = []
    for idx in idxes:
        if idx in same_idx_set:
            continue
        keep_idxes.append(idx)
        dist = distmat[idx]
        mask = dist < distance_thr
        mask[idx] = False
        same_image_idx = np.nonzero(mask)[0]
        if same_image_idx.size > 0:
            same_idx_set.update(same_image_idx.tolist())
    return keep_idxes


# 不包含start_thr和end_thr
def search_thr(keep_idxes, distmat, max_keep_count, dthr, start_thr, end_thr=None):
    cur_thr = distance_thr = start_thr
    while True:
        distance_thr += dthr
        if end_thr is not None:
            if distance_thr >= end_thr:
                break
        _keep_idxes = keep_idxes
        _keep_idxes = face_select(_keep_idxes, distance_thr, distmat)
        cur_thr = distance_thr
        keep_count = len(_keep_idxes)
        if keep_count <= max_keep_count:
            break

    start_distance_thr = distance_thr - dthr
    end_distance_thr = distance_thr
    return _keep_idxes, keep_count, start_distance_thr, end_distance_thr, cur_thr


def search_keep_idxes(keep_idxes, distmat, max_keep_count, start_thr, start_dthr, end_dthr):
    end_thr = None
    dthr = start_dthr
    while True:
        _keep_idxes, keep_count, start_thr, end_thr, cur_thr = search_thr(keep_idxes, distmat, max_keep_count, dthr, start_thr, end_thr)
        if (max_keep_count * 0.9 <= keep_count <= max_keep_count) or (dthr <= end_dthr):
            break
        else:
            dthr *= 0.5
    return _keep_idxes, cur_thr


def split_face_select(idxes, distance_thr, feats):
    same_idx_set = set()
    keep_idxes = []
    for idx in idxes:
        if idx in same_idx_set:
            continue
        keep_idxes.append(idx)
        distmat = cosine_distance2(feats, feats[idx:idx + 1])
        distmat = distmat.t().cpu().numpy()
        dist = distmat[0]
        mask = dist < distance_thr
        mask[idx] = False
        same_image_idx = np.nonzero(mask)[0]
        if same_image_idx.size > 0:
            same_idx_set.update(same_image_idx.tolist())
    return keep_idxes


# 不包含start_thr和end_thr
def split_search_thr(keep_idxes, feats, max_keep_count, dthr, start_thr, end_thr=None):
    cur_thr = distance_thr = start_thr
    while True:
        distance_thr += dthr
        if end_thr is not None:
            if distance_thr >= end_thr:
                break
        _keep_idxes = keep_idxes
        _keep_idxes = split_face_select(_keep_idxes, distance_thr, feats)
        cur_thr = distance_thr
        keep_count = len(_keep_idxes)
        if keep_count <= max_keep_count:
            break

    start_distance_thr = distance_thr - dthr
    end_distance_thr = distance_thr
    return _keep_idxes, keep_count, start_distance_thr, end_distance_thr, cur_thr


def split_search_keep_idxes(keep_idxes, feats, max_keep_count, start_thr, start_dthr, end_dthr):
    end_thr = None
    dthr = start_dthr
    while True:
        _keep_idxes, keep_count, start_thr, end_thr, cur_thr = split_search_thr(keep_idxes, feats, max_keep_count, dthr, start_thr, end_thr)
        if (max_keep_count * 0.9 <= keep_count <= max_keep_count) or (dthr <= end_dthr):
            break
        else:
            dthr *= 0.5
    return _keep_idxes, cur_thr


def auto_select_run(file_queue, out_queue, id, args, output_dir):
    DEVICE_ID = str(id % args.gpu_count)
    os.environ['CUDA_VISIBLE_DEVICES'] = DEVICE_ID

    DEVICE = 'cuda:0'

    import torch
    from torch.backends import cudnn

    cudnn.benchmark = True

    same_image_distance_thr = 1 - args.same_face_score_thr
    try:
        while True:
            file_info = file_queue.get(timeout=5)
            if file_info is None:
                break
            faceid, files, features = file_info
            total_image_count = len(files)
            if total_image_count == 1:
                out_queue.put(([], False, id))
                continue

            if not args.cpu:
                feats = features.cuda(device=DEVICE)
            else:
                feats = features
            feats.requires_grad = False
            # feats = torch.nn.functional.normalize(feats, dim=1, p=2)
            split_process = feats.size(0) > args.split_process_feature_count
            if split_process:
                dist_sum = cosine_distance_ex(feats)
            else:
                distmat = cosine_distance(feats)
                dist_sum = distmat.sum(axis=1)
                if not args.cpu:
                    distmat = distmat.cpu()
                    distmat = distmat.numpy()
            if not args.cpu:
                dist_sum = dist_sum.cpu()

            # distmat = cosine_distance(feats)
            # dist_sum = distmat.sum(axis=1)
            # if not args.cpu:
            #     distmat = distmat.cpu()
            #     dist_sum = dist_sum.cpu()
            # distmat = distmat.numpy()

            dist_sum = dist_sum.numpy()
            dist_sum_sort = dist_sum.argsort()
            # dist_sum_sort = dist_sum_sort[::-1]  # 倒序, 与其他图片差别大的排前面
            same_idx_set = set()
            keep_idxes = []
            # min_dist = 1.0
            for idx in dist_sum_sort:
                if idx in same_idx_set:
                    continue
                keep_idxes.append(idx)
                if split_process:
                    distmat = cosine_distance2(feats, feats[idx:idx+1])
                    distmat = distmat.t().cpu().numpy()
                    dist = distmat[0]
                else:
                    dist = distmat[idx]
                mask = dist < same_image_distance_thr
                mask[idx] = False
                # dist[idx] = 1.0
                # _min_dist = dist.min()
                # if min_dist > _min_dist:
                #     min_dist = _min_dist
                same_image_idx = np.nonzero(mask)[0]
                if same_image_idx.size > 0:
                    same_idx_set.update(same_image_idx.tolist())
            keep_count = len(keep_idxes)
            if keep_count < args.min_face_image_count:
                out_queue.put(([], False, id))
                continue
            if keep_count > args.max_face_image_count:
                # distance_thr = same_image_distance_thr if min_dist < same_image_distance_thr else min_dist
                distance_thr = same_image_distance_thr
                if split_process:
                    keep_idxes, cur_thr = split_search_keep_idxes(keep_idxes, feats, args.max_face_image_count, distance_thr, 0.1, 0.001)
                else:
                    keep_idxes, cur_thr = search_keep_idxes(keep_idxes, distmat, args.max_face_image_count, distance_thr, 0.1, 0.001)
                # print(len(keep_idxes), cur_thr)
            keep_files = [files[idx] for idx in keep_idxes]
            keep_feats = features[keep_idxes].clone()
            keep_feats.share_memory_()
            if args.show:
                show_files = []
                for keep_file in keep_files:
                    if 0 == args.mode:
                        keep_file = keep_file.strip().split()[-2]
                    else:
                        keep_file = keep_file.strip().split()[-1]
                    show_files.append(keep_file)
                if show_images(args.image_dir, show_files, total_image_count, wait=True):
                    break
            else:
                while out_queue.qsize() > 1000:
                    time.sleep(0.01)
                out_queue.put(([faceid, keep_files, keep_feats], False, id))
    except Exception as e:
        if str(e) != '':
            print(e)

    out_queue.put(([], True, id))


def out_run(out_queue, args, total_face_count, output_dir):
    face_count = 0
    feat_count = 0
    list_file_path = os.path.join(output_dir, 'FaceRecSel.txt')
    feature_file_path = os.path.join(output_dir, 'FaceRecFeatSel.dat')
    list_file = open(list_file_path, 'w')
    feature_file = open(feature_file_path, 'w')
    try:
        finish_worker_count = 0
        while True:
            file_info = out_queue.get(block=True)
            if file_info is None:
                break
            faceid_file_feats, finish, id = file_info
            if len(faceid_file_feats) > 0:
                face_count += 1
                faceid, files, features = faceid_file_feats
                feat_count += len(files)
                for file, feature in zip(files, features):
                    list_file.write(file)
                    feature = feature.numpy()
                    feature.tofile(feature_file)
            if face_count % 100 == 0:
                print('{:06f}, Proc{}, {}/{}, {}'.format(time.time(), id, face_count, total_face_count, feat_count))
            if finish:
                print('Proc{}, {}/{}, {} finish'.format(id, face_count, total_face_count, feat_count))
                finish_worker_count += 1
                if args.proccess_count <= finish_worker_count:
                    break
                continue
    except Exception as e:
        print(e)
    list_file.close()
    feature_file.close()
    os.system('chmod a+wr {}'.format(list_file_path))
    os.system('chmod a+wr {}'.format(feature_file_path))


def auto_select_mp(args, faceid_file_features):
    random.seed(123456)
    output_dir = args.output_dir

    file_queue = multiprocessing.Queue()
    out_queue = multiprocessing.Queue()

    total_face_count = len(faceid_file_features)

    workers = []
    for i in range(args.proccess_count):
        workers.append(multiprocessing.Process(target=auto_select_run, args=(file_queue, out_queue, i, args, output_dir)))

    out_worker = multiprocessing.Process(target=out_run, args=(out_queue, args, total_face_count, output_dir))

    for i in range(args.proccess_count):
        workers[i].start()

    out_worker.start()
    print('Start Time: {:06f}'.format(time.time()))
    wait_count_thr = args.proccess_count * 100
    for faceid, file_features in faceid_file_features.items():
        while file_queue.qsize() > wait_count_thr:
            time.sleep(0.01)
        files = file_features[0]
        features = file_features[1]
        file_count = len(files)
        if file_count > args.max_process_feature_count:
            idxes = list(range(file_count))
            random.shuffle(idxes)
            idxes = idxes[:args.max_process_feature_count]
            files = [files[idx] for idx in idxes]
            features = [features[idx] for idx in idxes]
        features = np.concatenate(features, axis=0)
        file_features[1] = None
        # features = multiprocessing.RawArray('f', features.ravel())
        features = torch.from_numpy(features)
        features.share_memory_()
        file_queue.put([faceid, files, features])

    out_worker.join()

    for i in range(args.proccess_count):
        workers[i].join()

    # for i, [seq_file_path, seq_name] in enumerate(zip(seq_files, seq_name_list)):
    #     if i == 0:
    #         continue
    #     file_queue.put([seq_file_path, seq_name])
    #     auto_select_run(file_queue, out_queue, 0, args, output_dir)
    #     break


def read_features(file_list, feature_file, feat_dim, mode=0):
    faceid_subfaceid_file_features = OrderedDict()
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            if mode == 0:
                # 015ff272-3768-4d4c-bebc-c14dc93065cd 477fc604-1479-4745-b8e9-9440a6ac8187 f8e4735b44724e25077106e0eaef15f0.jpg 0
                lines = line.strip().split()
                faceid = lines[0]
                subfaceid = lines[1]
                # file_name = lines[2]
                # mode = int(lines[3])
                faceid_subfaceid = '{}_{}'.format(faceid, subfaceid)
            elif mode == 1:
                # gd-zs-sqqjd-dxssjyeq-150100414a54443452064d2742c73600-video-20190913_225440/0001/0001_F0001B00.jpg
                lines = line.strip().split('/')
                faceid_subfaceid = '{}_{}'.format(lines[0], lines[1])
            elif mode == 2:
                # gd/zs/dqjd/xgly/150100414a5444345203bcd04294b500/video/20190912_154308/00121_000.jpg
                lines = line.strip().split('/')
                faceid_subfaceid = '{}_{}'.format(lines[-4], lines[-2])

            feature = np.fromfile(feat_file, dtype=np.float32, count=feat_dim).reshape(1, feat_dim)

            file_features = faceid_subfaceid_file_features.get(faceid_subfaceid, None)
            if file_features is None:
                files = [line]
                features = [feature]
                faceid_subfaceid_file_features[faceid_subfaceid] = [files, features]
            else:
                file_features[0].append(line)
                file_features[1].append(feature)

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return faceid_subfaceid_file_features


def main():
    parser = argparse.ArgumentParser(description="Select Face ReID Image")
    parser.add_argument(
        "--image_dir", default="/rootfs/media/kasim/Data/data/FaceImage/Image", help="path to image file list", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/Feature.FR.split0.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data/data/FaceImage/Feature.FR.split0.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data/data/FaceImage", help="path to output", type=str
    )
    parser.add_argument(
        "--feature_dim", default=256, help="feature dim", type=int
    )
    parser.add_argument('--proccess_count', type=int, default=8, help='detect proccess count')
    parser.add_argument('--gpu_count', type=int, default=2, help='detect gpu count')
    parser.add_argument("--show", action='store_true', help="show")
    parser.add_argument("--cpu", action='store_true', help="use cpu")
    parser.add_argument('--min_face_image_count', type=int, default=5, help='min face image count')
    parser.add_argument('--max_face_image_count', type=int, default=200, help='max face image count')
    parser.add_argument('--same_face_score_thr', type=float, default=0.95, help='same face score thr')
    parser.add_argument('--max_process_feature_count', type=int, default=4000, help='max process feature count')
    parser.add_argument('--split_process_feature_count', type=int, default=30000, help='split process feature count')
    parser.add_argument('--mode', type=int, default=0, help='file list mode')
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_subfaceid_file_features = read_features(args.file_list, args.feature_file, args.feature_dim, args.mode)

    auto_select_mp(args, faceid_subfaceid_file_features)

    print('finish!')


if __name__ == '__main__':
    main()
